from rest_framework_simplejwt.serializers import TokenObtainPairSerializer


class MyTokenObtainPairSerializer(TokenObtainPairSerializer):
    @classmethod
    def get_token(cls, user):
        token = super().get_token(user)
        token['user_id'] = user.id
        return token

    def validate(self, attrs):
        old_data = super().validate(attrs)
        data = {
            'code': 0,
            'msg': '登录成功',
            'data': {
                'user_id': self.user.id,
                'refresh': old_data['refresh'],
                'access': old_data['access'],
            }
        }
        return data
